import torch


def preprocess(conversations, tokenizer, max_tokens=None):
    """
    Preprocess the data by tokenizing.
    """

    all_input_ids = []  # 存储所有处理后的输入ID
    all_labels = []  # 存储所有的标签

    for conv in conversations:  # 对于每一组对话
        roles = [msg["role"] for msg in conv]  # 获取对话中每个人的角色，例如“SYSTEM”, “ASSISTANT”或“USER”
        messages = [msg["content"] for msg in conv]  # 获取对话中每个人的消息内容

        # 断言第一个角色不是“ASSISTANT”和最后一个角色是“ASSISTANT”
        #  这个可以使用也可以不使用
        assert roles[0] != "ASSISTANT"
        assert roles[-1] == "ASSISTANT"

        input_messages = []  # 存储需要输入的消息

        # 根据角色将消息添加到input_messages中，"ASSISTANT"和"USER"的消息都被添加
        for role, msg in zip(roles, messages):
            if role == "ASSISTANT":
                input_messages.append(msg)
            elif role == "USER":
                input_messages.append(msg)

        #使用ChatGLM3的tokeninzer进行token处理
        tokenized_input = tokenizer(input_messages, add_special_tokens=False)  # 对输入消息进行token化

        input_ids = []  # 初始化本次对话的输入ID
        labels = []  # 初始化本次对话的标签

        # 根据第一个角色是"SYSTEM"还是其他角色来添加初始的输入ID和标签
        if roles[0] == "SYSTEM":
            input_ids.extend([64790, 64792, 64794, 30910, 13])  #起始位置拼接特定的token ID
            input_ids.extend(tokenized_input.input_ids[0])
            labels.extend([-100] * (len(tokenized_input.input_ids[0]) + 5))     #将label设置成-100，这是由于在交叉熵计算时，-100对应的位置不参与损失值计算
        else:
            input_ids.extend([64790, 64792])    #起始位置拼接特定的token ID
            labels.extend([-100] * 2)    #将label设置成-100，这是由于在交叉熵计算时，-100对应的位置不参与损失值计算

            # 根据每个人的角色和token化的消息，添加输入ID和标签
        for role, msg in zip(roles, tokenized_input.input_ids):
            if role == "USER":
                if roles[0] == "SYSTEM":
                    labels.extend([-100] * (len(msg) + 5))
                    input_ids.extend([13, 64795, 30910, 13])
                else:
                    labels.extend([-100] * (len(msg) + 4))  #将label设置成-100，这里是USER提问部分，不参与损失函数计算
                    input_ids.extend([64795, 30910, 13])    #添加USER对话开始的起始符
                input_ids.extend(msg)   #将当前的消息token添加到输入ID列表中
                input_ids.extend([64796])   #添加USER对话结束符
                print("USER", msg)  # 打印USER的消息和对应的token IDs


            elif role == "ASSISTANT":  # 当角色为"ASSISTANT"时
                msg += [tokenizer.eos_token_id]  # 在消息后面添加一个结束token的ID
                #这里的作用
                labels.extend([30910, 13])  #添加ASSISTANT对话开始的起始符
                labels.extend(msg)  # 将当前的消息token添加到标签列表中

                input_ids.extend([30910, 13])  #添加ASSISTANT对话开始的起始符
                input_ids.extend(msg)  # 将当前的消息token添加到输入ID列表中
                print("ASSISTANT", msg)  # 打印ASSISTANT的消息和对应的token IDs

        if max_tokens is None:  # 如果没有设定最大token数量
            max_tokens = tokenizer.model_max_length  # 则使用tokenizer的模型最大长度作为最大token数量

        input_ids = torch.LongTensor(input_ids)[:max_tokens]  # 将输入ID列表转化为LongTensor，并截取前max_tokens个token
        labels = torch.LongTensor(labels)[:max_tokens]  # 将标签列表转化为LongTensor，并截取前max_tokens个token

        assert input_ids.shape == labels.shape  # 判断输入ID的tensor和标签的tensor形状是否相同，确保一一对应

        all_input_ids.append(input_ids)  # 将处理后的输入ID添加到所有输入ID列表中
        all_labels.append(labels)  # 将处理后的标签添加到所有标签列表中

    return dict(input_ids=all_input_ids, labels=all_labels)
